# examples/baostock_example.py
from sheng.stock.basic_data_v1.data_manager import StockDataManager
import pandas as pd

def baostock_demo():
    """Baostock数据源使用演示"""
    manager = StockDataManager()

    print("=== Baostock数据获取演示 ===")

    # 1. 获取股票列表
    print("1. 获取股票列表...")
    manager.update_stock_list('baostock')

    # 2. 更新日线数据（测试几只股票）
    test_symbols = ['000001', '600036', '601318', '000858']
    print("2. 更新日线数据...")
    manager.update_daily_data(test_symbols, 'baostock')

    # 3. 更新指标数据
    print("3. 更新指标数据...")
    manager.update_baostock_indicators(test_symbols)

    # 4. 更新分红数据
    print("4. 更新分红数据...")
    manager.update_dividend_data(test_symbols, '2023')

    # 5. 查询数据
    print("5. 查询数据...")
    df = manager.get_data('000001', '2023-01-01', '2023-12-31')

    print(f"获取到平安银行 {len(df)} 条日线数据")
    if not df.empty:
        print("最新5条数据:")
        print(df[['trade_date', 'open_price', 'close_price', 'change_rate', 'volume']].tail())

    # 6. 查询指标数据
    conn = manager.db._get_connection()
    indicators_df = pd.read_sql(
        "SELECT * FROM baostock_indicators WHERE symbol = '000001' ORDER BY trade_date DESC LIMIT 5",
        conn
    )
    conn.close()

    if not indicators_df.empty:
        print("\n最新5条指标数据:")
        print(indicators_df[['trade_date', 'pe_ratio', 'pb_ratio']])

def compare_data_sources():
    """比较不同数据源的数据"""
    manager = StockDataManager()

    symbols = ['000001', '600036']

    for symbol in symbols:
        print(f"\n=== 比较 {symbol} 的数据 ===")

        # 从不同数据源获取数据
        sources = ['baostock', 'akshare']

        for source in sources:
            try:
                # 这里需要根据数据源特性调整获取方式
                if source == 'baostock':
                    df = manager.get_data(symbol, '2023-12-01', '2023-12-31')
                    print(f"{source}: {len(df)} 条记录")
                elif source == 'akshare':
                    df = manager.get_data(symbol, '2023-12-01', '2023-12-31')
                    print(f"{source}: {len(df)} 条记录")
            except Exception as e:
                print(f"{source} 获取失败: {e}")

if __name__ == "__main__":
    baostock_demo()
    # compare_data_sources()